import logging
import time
from functools import wraps
from threading import local
from typing import Literal

# Thread local storage
_local = local()

# Define ANSI color codes
class Colors:
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    MAGENTA = '\033[95m'
    CYAN = '\033[96m'
    WHITE = '\033[97m'
    RESET = '\033[0m'

# Global variable to store averages
average_times = {}
display_average_times = {}
def get_average_time():
    return display_average_times 

def timing_decorator(
    verbose: Literal['all', 'self_only', 'child_only', 'none', 'force'] = 'all',
    unit: Literal['ms', 's'] = 'ms',
    alpha: float = 0.05 , # Smoothing factor for exponential moving average
    show_in_progress: bool = True
):
    def format_time(execution_time: float, func_name: str, avg_time: float) -> None:
        """Format and log execution time"""
        # Color the function name
        colored_func_name = f"{Colors.CYAN}{func_name}{Colors.RESET}"
        
        if unit == 'ms':
            exec_time_str = f"{execution_time * 1000:.1f} ms"
            avg_time_str = f"{avg_time * 1000:.1f} ms"
            colored_exec_time = f"{Colors.GREEN}{exec_time_str}{Colors.RESET}"
            colored_avg_time = f"{Colors.YELLOW}{avg_time_str}{Colors.RESET}"
            logging.info(f"{colored_func_name} took {colored_exec_time} to execute | Avg: {colored_avg_time}")
        elif unit == 's':
            exec_time_str = f"{execution_time:.2f} s"
            avg_time_str = f"{avg_time:.2f} s"
            colored_exec_time = f"{Colors.GREEN}{exec_time_str}{Colors.RESET}"
            colored_avg_time = f"{Colors.YELLOW}{avg_time_str}{Colors.RESET}"
            logging.info(f"{colored_func_name} took {colored_exec_time} to execute | Avg: {colored_avg_time}")
        else:
            logging.warning(f"Unknown unit '{unit}', using seconds by default: {execution_time:.2f} s")

    def measure_time(func, *args, **kwargs):
        """Measure function execution time"""
        try:
            start_time = time.time()
            result = func(*args, **kwargs)
            end_time = time.time()
            execution_time = end_time - start_time
            
            # Calculate average time using exponential moving average
            if func.__qualname__ not in average_times:
                average_times[func.__qualname__] = execution_time
            else:
                average_times[func.__qualname__] = (alpha * execution_time) + ((1 - alpha) * average_times[func.__qualname__])
            if show_in_progress is False:
            
                format_time(execution_time, func.__qualname__, average_times[func.__qualname__])
            else:
                display_average_times[str(func.__qualname__).split('.')[-1]] = f'{average_times[func.__qualname__] * 1000:.2f} ms' if unit == 'ms' else f'{average_times[func.__qualname__] :.2f} ms'
            return result
        except Exception as e:
            logging.error(f"Error occurred in {func.__qualname__}: {e}")
            raise

    def decorator(func):
        if verbose == 'child_only':
            return func
            
        if verbose == 'none':
            @wraps(func)
            def wrapper(*args, **kwargs):
                _local.in_timing = True
                return func(*args, **kwargs)
            return wrapper

        @wraps(func)
        def wrapper(*args, **kwargs):
            # If already in timing context and not in 'force' mode, execute directly
            if hasattr(_local, 'in_timing') and verbose != 'force':
                return func(*args, **kwargs)
            
            if verbose == 'self_only':
                try:
                    _local.in_timing = True
                    return measure_time(func, *args, **kwargs)
                finally:
                    delattr(_local, 'in_timing')
            
            return measure_time(func, *args, **kwargs)
            
        return wrapper
    
    return decorator


if __name__ == '__main__':
    if not logging.getLogger().handlers:
        logging.basicConfig(
            level=logging.DEBUG,
            format='%(asctime)s %(filename)s %(levelname)s %(message)s',
            datefmt='%a, %d %b %Y %H:%M:%S',
            # filename='myapp.log',
            # filemode='w',
        )
    # Test
    @timing_decorator()
    def inner_function():
        time.sleep(0.1)
        return "inner done"

    @timing_decorator(verbose='all')
    def outer_function_all():
        time.sleep(0.2)
        inner_function()
        return "outer done"
      
    @timing_decorator(verbose='self_only')
    def outer_function_self_only():
        time.sleep(0.2)
        inner_function()
        return "outer done"    
    
    @timing_decorator(verbose='child_only')
    def outer_function_child_only():
        time.sleep(0.2)
        inner_function()
        return "outer done"    
    
    @timing_decorator(verbose='none')
    def outer_function_none():
        time.sleep(0.2)
        inner_function()
        return "outer done"
    
    inner_function()

    print('--------------------')
    print('--------all---------')
    print('--------------------')
    outer_function_all()

    print('--------------------')
    print('-----self_only------')
    print('--------------------')
    outer_function_self_only()

    print('--------------------')
    print('-----child_only-----')
    print('--------------------')
    outer_function_child_only()

    print('--------------------')
    print('--------none--------')
    print('--------------------')
    outer_function_none()